from dotenv import load_dotenv, find_dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings

_ = load_dotenv(find_dotenv())
embeddings  = ZhipuAIEmbeddings(model="embedding-3")

text = "你变瘦了。"
text_list = [
    "你减肥成功了。",
    "你变苗条了。",
    "你又胖了。",
    "你好瘦啊！"
]

query_result = embeddings.embed_query(text)
# print(f"len(query_result)={len(query_result)}")
#print(f"query_result={query_result}")
doc_result = embeddings.embed_documents(text_list)
# print(f"len(doc_result)={len(doc_result)}")
# print(f"len(doc_result[0])={len(doc_result[0])}")

import numpy as np
def cosine_similarity(vec1, vec2):
    dot_product = np.dot(vec1, vec2)
    norm_vec1 = np.linalg.norm(vec1)
    norm_vec2 = np.linalg.norm(vec2)
    return dot_product / (norm_vec1 * norm_vec2)

for v, t in zip(doc_result, text_list) :
    similarity = cosine_similarity(query_result, v)
    print(f"{t}\t:{similarity}") 
